{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 🧪 ADK Application Testing\n",
    "\n",
    "This notebook demonstrates how to test an ADK (Agent Development Kit) application.\n",
    "It covers both local and remote testing, both with Agent Engine and Cloud Run.\n",
    "\n",
    "> **Note**: This notebook assumes that the agent files are stored in the `app` folder. If your agent files are located in a different directory, please update all relevant file paths and references accordingly."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Set Up Your Environment\n",
    "\n",
    "> **Note:** For best results, use the same `.venv` created for local development with `uv` to ensure dependency compatibility and avoid environment-related issues."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Uncomment the following lines if you're not using the virtual environment created by uv\n",
    "# import sys\n",
    "\n",
    "# sys.path.append(\"../\")\n",
    "# !pip install google-cloud-aiplatform a2a-sdk --upgrade"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Import libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "import requests\n",
    "import vertexai"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Initialize Vertex AI Client"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize the Vertex AI client\n",
    "LOCATION = \"us-central1\"\n",
    "\n",
    "client = vertexai.Client(\n",
    "    location=LOCATION,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## If you are using Agent Engine\n",
    "See more documentation at [Agent Engine Overview](https://cloud.google.com/vertex-ai/generative-ai/docs/agent-engine/overview)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Remote Testing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set to None to auto-detect from ./deployment_metadata.json, or specify manually\n",
    "# \"projects/PROJECT_ID/locations/us-central1/reasoningEngines/ENGINE_ID\"\n",
    "REASONING_ENGINE_ID = None\n",
    "\n",
    "if REASONING_ENGINE_ID is None:\n",
    "    try:\n",
    "        with open(\"../deployment_metadata.json\") as f:\n",
    "            metadata = json.load(f)\n",
    "            REASONING_ENGINE_ID = metadata.get(\"remote_agent_engine_id\")\n",
    "    except (FileNotFoundError, json.JSONDecodeError):\n",
    "        pass\n",
    "\n",
    "print(f\"Using REASONING_ENGINE_ID: {REASONING_ENGINE_ID}\")\n",
    "# Get the existing agent engine\n",
    "remote_agent_engine = client.agent_engines.get(name=REASONING_ENGINE_ID)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "async for event in remote_agent_engine.async_stream_query(\n",
    "    message=\"hi!\", user_id=\"test\"\n",
    "):\n",
    "    print(event)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": "remote_agent_engine.register_feedback(\n    feedback={\n        \"score\": 5,\n        \"text\": \"Great response!\",\n        \"user_id\": \"test-user-123\",\n        \"session_id\": \"test-session-123\",\n    }\n)"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Local Testing\n",
    "\n",
    "You can import directly the AgentEngineApp class within your environment. \n",
    "To run the agent locally, follow these steps:\n",
    "1. Make sure all required packages are installed in your environment\n",
    "2. The recommended approach is to use the same virtual environment created by the 'uv' tool\n",
    "3. You can set up this environment by running 'make install' from your agent's root directory\n",
    "4. Then select this kernel (.venv folder in your project) in your Jupyter notebook to ensure all dependencies are available"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from app.agent_engine_app import agent_engine\n",
    "\n",
    "agent_engine.set_up()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "async for event in agent_engine.async_stream_query(message=\"hi!\", user_id=\"test\"):\n",
    "    print(event)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## If you are using Cloud Run"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Remote Testing\n",
    "\n",
    "For more information about authenticating HTTPS requests to Cloud Run services, see:\n",
    "[Cloud Run Authentication Documentation](https://cloud.google.com/run/docs/triggering/https-request)\n",
    "\n",
    "Remote testing involves using a deployed service URL instead of localhost.\n",
    "\n",
    "Authentication is handled using GCP identity tokens instead of local credentials."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ID_TOKEN = get_ipython().getoutput(\"gcloud auth print-identity-token -q\")[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "SERVICE_URL = \"YOUR_SERVICE_URL_HERE\"  # Replace with your Cloud Run service URL"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You'll need to first create a Session"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "user_id = \"test_user_123\"\n",
    "session_data = {\"state\": {\"preferred_language\": \"English\", \"visit_count\": 1}}\n",
    "\n",
    "session_url = f\"{SERVICE_URL}/apps/app/users/{user_id}/sessions\"\n",
    "headers = {\"Content-Type\": \"application/json\", \"Authorization\": f\"Bearer {ID_TOKEN}\"}\n",
    "\n",
    "session_response = requests.post(session_url, headers=headers, json=session_data)\n",
    "print(f\"Session creation status code: {session_response.status_code}\")\n",
    "print(f\"Session creation response: {session_response.json()}\")\n",
    "session_id = session_response.json()[\"id\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then you will be able to send a message"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "message_data = {\n",
    "    \"app_name\": \"app\",\n",
    "    \"user_id\": user_id,\n",
    "    \"session_id\": session_id,\n",
    "    \"new_message\": {\"role\": \"user\", \"parts\": [{\"text\": \"Hello! Weather in New york?\"}]},\n",
    "    \"streaming\": True,\n",
    "}\n",
    "\n",
    "message_url = f\"{SERVICE_URL}/run_sse\"\n",
    "message_response = requests.post(\n",
    "    message_url, headers=headers, json=message_data, stream=True\n",
    ")\n",
    "\n",
    "print(f\"Message send status code: {message_response.status_code}\")\n",
    "\n",
    "# Print streamed response\n",
    "for line in message_response.iter_lines():\n",
    "    if line:\n",
    "        line_str = line.decode(\"utf-8\")\n",
    "        if line_str.startswith(\"data: \"):\n",
    "            event_json = line_str[6:]\n",
    "            event = json.loads(event_json)\n",
    "            print(f\"Received event: {event}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Local Testing\n",
    "\n",
    "> You can run the application locally via the `make local-backend` command."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Create a session\n",
    " Create a new session with user preferences and state information\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "user_id = \"test_user_123\"\n",
    "session_data = {\"state\": {\"preferred_language\": \"English\", \"visit_count\": 1}}\n",
    "\n",
    "session_url = f\"http://127.0.0.1:8000/apps/app/users/{user_id}/sessions\"\n",
    "headers = {\"Content-Type\": \"application/json\"}\n",
    "\n",
    "session_response = requests.post(session_url, headers=headers, json=session_data)\n",
    "print(f\"Session creation status code: {session_response.status_code}\")\n",
    "print(f\"Session creation response: {session_response.json()}\")\n",
    "session_id = session_response.json()[\"id\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Send a message\n",
    "Send a message to the backend service and receive a streaming response\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "message_data = {\n",
    "    \"app_name\": \"app\",\n",
    "    \"user_id\": user_id,\n",
    "    \"session_id\": session_id,\n",
    "    \"new_message\": {\"role\": \"user\", \"parts\": [{\"text\": \"Hello! Weather in New york?\"}]},\n",
    "    \"streaming\": True,\n",
    "}\n",
    "\n",
    "message_url = \"http://127.0.0.1:8000/run_sse\"\n",
    "message_response = requests.post(\n",
    "    message_url, headers=headers, json=message_data, stream=True\n",
    ")\n",
    "\n",
    "print(f\"Message send status code: {message_response.status_code}\")\n",
    "\n",
    "# Print streamed response\n",
    "for line in message_response.iter_lines():\n",
    "    if line:\n",
    "        line_str = line.decode(\"utf-8\")\n",
    "        if line_str.startswith(\"data: \"):\n",
    "            event_json = line_str[6:]\n",
    "            event = json.loads(event_json)\n",
    "            print(f\"Received event: {event}\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "myagent-1762384391",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}